#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/epoll.h> // epoll functions
#include <fcntl.h>     // nonblocking
#include <stdlib.h>
#include <errno.h>
#include <stdio.h>
#include <string.h>

#define MAX_EVENTS 10
#define MAX_LINE 1024
#define LISTENQ 1024
#define SERVER_PORT 6888

// Function to handle data from a client connection
int handle(int connfd)
{
    int nread;
    char buf[MAX_LINE];
    memset(buf, 0, MAX_LINE);

    nread = read(connfd, buf, MAX_LINE - 1);

    if (nread == 0)
    {
        printf("Client fd %d closed the connection\n", connfd);
        return -1;
    }
    if (nread < 0)
    {
        if (errno == EAGAIN || errno == EWOULDBLOCK)
        {
            printf("Read would block on fd %d\n", connfd);
            return 0;
        }
        else if (errno == ECONNRESET)
        {
            printf("Client fd %d connection reset.\n", connfd);
            return -1;
        }
        else
        {
            perror("read error");
            return -1;
        }
    }

    buf[nread] = '\0';
    printf("Received from fd %d: %s\n", connfd, buf);
    return 0;
}

// Function to set a socket to non-blocking mode
int setnonblocking(int sockfd)
{
    if (fcntl(sockfd, F_SETFL, fcntl(sockfd, F_GETFL, 0) | O_NONBLOCK) == -1)
    {
        perror("fcntl set nonblocking error");
        return -1;
    }
    return 0;
}

int main(int argc, char **argv)
{
    int listenfd, connfd, epollfd, nfds;
    struct sockaddr_in servaddr, cliaddr;
    socklen_t socklen = sizeof(struct sockaddr_in);
    struct epoll_event ev, events[MAX_EVENTS];
    char buf[MAX_LINE];
    int acceptCount = 0;

    listenfd = socket(AF_INET, SOCK_STREAM, 0);
    if (listenfd == -1)
    {
        perror("create socket failed");
        return -1;
    }

    int opt = 1;
    if (setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0)
    {
        perror("setsockopt reuseaddr error");
        close(listenfd);
        return -1;
    }

    if (setnonblocking(listenfd) < 0)
    {
        close(listenfd);
        return -1;
    }

    memset(&servaddr, 0, sizeof(servaddr));
    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = INADDR_ANY;
    servaddr.sin_port = htons(SERVER_PORT);

    if (bind(listenfd, (struct sockaddr *)&servaddr, sizeof(struct sockaddr)) == -1)
    {
        perror("bind error");
        close(listenfd);
        return -1;
    }

    if (listen(listenfd, LISTENQ) == -1)
    {
        perror("listen error");
        close(listenfd);
        return -1;
    }

    epollfd = epoll_create1(0);
    if (epollfd == -1)
    {
        perror("epoll_create1 failed");
        close(listenfd);
        return -1;
    }

    ev.events = EPOLLIN; // Use LT mode for simplicity matching other examples
    // ev.events = EPOLLIN | EPOLLET; // Use ET for higher performance if handle() logic is adjusted
    ev.data.fd = listenfd;
    if (epoll_ctl(epollfd, EPOLL_CTL_ADD, listenfd, &ev) < 0)
    {
        fprintf(stderr, "epoll_ctl add listenfd error: fd=%d\n", listenfd);
        close(listenfd);
        close(epollfd);
        return -1;
    }

    printf("epoll server startup. port: %d\n", SERVER_PORT);

    for (;;)
    {
        nfds = epoll_wait(epollfd, events, MAX_EVENTS, -1);
        if (nfds == -1)
        {
            if (errno == EINTR)
                continue;
            perror("epoll_wait error");
            continue;
        }

        for (int n = 0; n < nfds; ++n)
        {
            int current_fd = events[n].data.fd;

            if (current_fd == listenfd)
            {
                // Accept loop for ET, single accept usually sufficient for LT
                while (1)
                {
                    connfd = accept(listenfd, (struct sockaddr *)&cliaddr, &socklen);
                    if (connfd < 0)
                    {
                        if (errno == EAGAIN || errno == EWOULDBLOCK)
                        {
                            break; // No more connections pending
                        }
                        else if (errno == EINTR)
                        {
                            continue; // Interrupted, retry accept
                        }
                        else
                        {
                            perror("accept error");
                            break;
                        }
                    }

                    sprintf(buf, "accept form %s:%d\n", inet_ntoa(cliaddr.sin_addr), ntohs(cliaddr.sin_port));
                    printf("Connection %d: %s", ++acceptCount, buf);

                    if (setnonblocking(connfd) < 0)
                    {
                        close(connfd);
                        continue;
                    }

                    ev.events = EPOLLIN; // LT mode
                    // ev.events = EPOLLIN | EPOLLET; // ET mode
                    ev.data.fd = connfd;
                    if (epoll_ctl(epollfd, EPOLL_CTL_ADD, connfd, &ev) < 0)
                    {
                        perror("epoll_ctl add connfd error");
                        close(connfd);
                    }
                    // If using LT, break after one accept
                    break; // Remove this break if using ET and want to accept all pending
                }
            }
            else
            {
                // Handle client data or error
                if ((events[n].events & EPOLLERR) || (events[n].events & EPOLLHUP) || !(events[n].events & EPOLLIN))
                {
                    fprintf(stderr, "epoll error/hup on fd %d\n", current_fd);
                    epoll_ctl(epollfd, EPOLL_CTL_DEL, current_fd, NULL);
                    close(current_fd);
                }
                else if (events[n].events & EPOLLIN)
                {
                    if (handle(current_fd) < 0)
                    {
                        // Client closed or error in handle()
                        epoll_ctl(epollfd, EPOLL_CTL_DEL, current_fd, NULL);
                        close(current_fd); // Close socket here after removing from epoll
                    }
                }
            }
        }
    }

    close(listenfd);
    close(epollfd);
    return 0;
}